// Copyright 2014 The mqrouter Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package xkafka

import (
	"context"
	"encoding/json"
	"gitee.com/go-mid/infra/xmq/mqconf"
	"time"

	"github.com/segmentio/kafka-go"

	"gitee.com/go-mid/infra/xtime"
	"gitee.com/go-mid/infra/xtrace"
)

type ReaderConfig struct {
	mqconf.MQConfig
	kafka.ReaderConfig
}
type KafkaReader struct {
	*kafka.Reader
	config kafka.ReaderConfig
	stop   chan struct{}
}

func NewKafkaReader(readerConfig *ReaderConfig) *KafkaReader {
	reader := kafka.NewReader(readerConfig.ReaderConfig)

	kafkaReader := &KafkaReader{
		Reader: reader,
		stop:   make(chan struct{}),
		config: readerConfig.ReaderConfig,
	}
	go kafkaReader.init()
	return kafkaReader
}

func (m *KafkaReader) init() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			stats := m.Stats()
			statKafkaReader(m.Config().Topic, &stats)
		case <-m.stop:
			return
		}
	}
}
func (m *KafkaReader) setKafkaSpanTags(span xtrace.Span) {
	config := m.Config()
	span.SetTag(xtrace.TagSpanKind, xtrace.SpanKindConsumer)
	span.SetTag(xtrace.TagMessageBusDestination, config.Topic)
	span.SetTag(xtrace.TagMessagingSystem, xtrace.MessagingSystemKafka)
	span.SetTag(xtrace.TagMessagingDestinationKind, xtrace.MessagingDestinationKindTopic)
	span.SetTag(xtrace.TagMessagingDestination, config.Topic)

	if config.GroupID != "" {
		span.SetTag(xtrace.TagPalfishKafkaConsumerGroupID, config.GroupID)
		span.SetTag(xtrace.TagMessagingKafkaConsumerGroup, config.GroupID)
	}

	if config.Partition != 0 {
		span.SetTag(xtrace.TagPalfishKafkaConsumerPartition, config.Partition)
		span.SetTag(xtrace.TagMessagingKafkaPartition, config.Partition)
	}
}

func (m *KafkaReader) ReadMsg(ctx context.Context, v interface{}) error {
	span := xtrace.SpanFromContext(ctx)
	if span != nil {
		m.setKafkaSpanTags(span)
	}

	st := xtime.NewTimeStat()
	msg, err := m.ReadMessage(ctx)
	mqconf.StatReqDuration(ctx, m.Config().Topic, "KafkaReader.ReadMessage", mqconf.TraceMessageBusTypeKafka, st.Millisecond())

	if err != nil {
		return err
	}

	err = json.Unmarshal(msg.Value, v)
	if err != nil {
		return err
	}

	return nil
}

func (m *KafkaReader) FetchMsg(ctx context.Context, v interface{}) (mqconf.AckHandler, error) {
	span := xtrace.SpanFromContext(ctx)
	if span != nil {
		m.setKafkaSpanTags(span)
	}

	st := xtime.NewTimeStat()
	msg, err := m.FetchMessage(ctx)
	mqconf.StatReqDuration(ctx, m.Config().Topic, "KafkaReader.FetchMessage", mqconf.TraceMessageBusTypeKafka, st.Millisecond())

	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(msg.Value, v)
	if err != nil {
		return nil, err
	}

	return &KafkaHandler{
		msg:    msg,
		reader: m.Reader,
	}, nil
}

func (m *KafkaReader) Close(ctx context.Context) error {
	m.stop <- struct{}{}
	return m.Reader.Close()
}

func (m *KafkaReader) SetOffsetAt(ctx context.Context, t time.Time) error {
	return m.Reader.SetOffsetAt(ctx, t)
}

func (m *KafkaReader) SetOffset(ctx context.Context, offset int64) error {
	return m.Reader.SetOffset(offset)
}

type KafkaHandler struct {
	msg    kafka.Message
	reader *kafka.Reader
}

func (m *KafkaHandler) Ack(ctx context.Context) error {
	return m.reader.CommitMessages(ctx, m.msg)
}
